from collections.abc import Sequence
from typing import NamedTuple

from zkl_aiutils_datasets import Dataset, DatasetIterator, SkipExhausted


class SlidingWindowSample(NamedTuple):
    data: Sequence
    head: int
    tail: int


class SlidingWindowIterator(DatasetIterator[SlidingWindowSample]):
    def __init__(self,
        sequence: Sequence,
        win_size: int,
        win_stride: int,
        state: int | None = None,
    ):
        self.sequence = sequence
        self.win_size = win_size
        self.win_stride = win_stride
        self.head = state if state is not None else 0

    def __next__(self) -> SlidingWindowSample:
        head = self.head
        tail = head + self.win_size
        if tail > len(self.sequence):
            raise StopIteration
        self.head += self.win_stride
        return SlidingWindowSample(self.sequence[head:tail], head, tail)

    def __skip__(self, samples_n: int):
        if samples_n < 0:
            raise ValueError(f"Expected samples_n>=0, got {samples_n=}.")
        if samples_n == 0:
            return
        remain_samples_n = (len(self.sequence) - self.head - self.win_size) // self.win_stride + 1 \
            if len(self.sequence) - self.head >= self.win_size else 0
        if samples_n > remain_samples_n:
            self.head += remain_samples_n * self.win_stride
            raise SkipExhausted(remain_samples_n)
        self.head += samples_n * self.win_stride

    def __pause__(self) -> int:
        return self.head


class SlidingWindowDataset(Dataset[SlidingWindowSample]):
    def __init__(self, sequence: Sequence, win_size: int, win_stride: int):
        self.sequence = sequence
        self.win_size = win_size
        self.win_stride = win_stride

    def __len__(self) -> int:
        return (len(self.sequence) - self.win_size) // self.win_stride + 1 \
            if len(self.sequence) >= self.win_size else 0

    def __getitem__(self, index: int, /) -> SlidingWindowSample:
        head = index * self.win_stride
        tail = head + self.win_size
        return SlidingWindowSample(self.sequence[head:tail], head, tail)

    def __iter__(self) -> SlidingWindowIterator:
        return SlidingWindowIterator(self.sequence, self.win_size, self.win_stride)

    def __resume__(self, state: int | None = None, /) -> SlidingWindowIterator:
        return SlidingWindowIterator(self.sequence, self.win_size, self.win_stride, state)
